package de.distMLP.communication;

import java.io.IOException;

import de.distMLP.messages.DoubleMessage;
import de.distMLP.messages.HamaMessage;
import de.distMLP.messages.VectorMessage;
import de.distMLP.messages.WeightMatrixMessage;
import de.mlp_distributed.mlp.math.mahout.Matrix;
import de.mlp_distributed.mlp.math.mahout.Vector;

public class MessageAggregator {

	protected HamaMessage aggregateMessage(final HamaMessage a, final HamaMessage b) throws IOException {
		HamaMessage aggregateMessage = null;

		if ((a instanceof DoubleMessage) && (b instanceof DoubleMessage)) {
			aggregateMessage = this.handleDoubleMessage(a, b);
		} else if ((a instanceof WeightMatrixMessage) && (b instanceof WeightMatrixMessage)) {
			aggregateMessage = this.handleWeightmatrixMessage(a, b);
		} else if ((a instanceof VectorMessage) && (b instanceof VectorMessage)) {
			final Vector sum = MessageAggregator.addVector(((VectorMessage) a).getVector(), ((VectorMessage) b).getVector());
			aggregateMessage = new VectorMessage(sum);
		}
		return aggregateMessage;
	}

	public HamaMessage handleDoubleMessage(final HamaMessage a, final HamaMessage b) {
		HamaMessage aggregateMessage;
		aggregateMessage = new DoubleMessage(MessageAggregator.addDouble(((DoubleMessage) a).getDoubleValue(),
				((DoubleMessage) b).getDoubleValue()));
		return aggregateMessage;
	}

	/**
	 * sums the matrices, cost and checks if stop training was send.
	 * 
	 * @param a
	 * @param b
	 * @return
	 * @throws IOException
	 */
	public HamaMessage handleWeightmatrixMessage(final HamaMessage a, final HamaMessage b) throws IOException {
		HamaMessage aggregateMessage;
		final Matrix[] sum = MessageAggregator.addMatrixArray(((WeightMatrixMessage) a).getMatrix(), ((WeightMatrixMessage) b).getMatrix());
		final double cost = MessageAggregator.addDouble(((WeightMatrixMessage) a).getCost(), ((WeightMatrixMessage) b).getCost());
		boolean stopTraining = false;
		if (((WeightMatrixMessage) a).isStopTraining() || ((WeightMatrixMessage) b).isStopTraining()) {
			stopTraining = true;
		}
		aggregateMessage = new WeightMatrixMessage(sum, stopTraining, cost);
		return aggregateMessage;
	}

	/**
	 * 
	 * @param a
	 * @param b
	 * @return new matrix array containing a + b
	 */
	protected static Matrix[] addMatrixArray(final Matrix[] a, final Matrix[] b) {
		final Matrix[] tmpMatrix = new Matrix[a.length];

		for (int i = 0; i < a.length; i++) {
			if (a[i] == null) {
				tmpMatrix[i] = b[i];
			} else {
				tmpMatrix[i] = a[i].plus(b[i]);
			}
		}
		return tmpMatrix;
	}

	protected static void divideMatrixArrayByDouble(final Matrix[] weightChanges, final double value) {
		for (int i = 0; i < weightChanges.length; i++) {
			weightChanges[i] = weightChanges[i].divide(value);
		}
	}

	protected static Vector addVector(final Vector a, final Vector b) {
		return a.plus(b);
	}

	protected static double addDouble(final double a, final double b) {
		return a + b;
	}
}
